{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Notebook for running a DM-BCA model\n",
    "\n",
    "This notebook can be used to run a BCA model pretrained on the discourse marker prediction task, and trained on TOEFL LDC data (bca_dm_toefl_model). Please use the script TOEFL_dataParse.py to preprocess the TOEFL data. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import print_function, division\n",
    "\n",
    "import os\n",
    "import os.path\n",
    "import pandas as pd\n",
    "from io import StringIO\n",
    "import io\n",
    "import unicodedata\n",
    "import re\n",
    "import random\n",
    "\n",
    "import tensorflow as tf\n",
    "import numpy as np\n",
    "np.set_printoptions(threshold = 10000)\n",
    "import collections\n",
    "import random\n",
    "\n",
    "from tensorflow.contrib.rnn import LSTMCell as Cell #for GRU: custom implementation with normalization\n",
    "from tensorflow.python.ops.rnn import dynamic_rnn as rnn\n",
    "from tensorflow.python.ops.rnn import bidirectional_dynamic_rnn as bi_rnn\n",
    "from tensorflow.contrib.rnn import DropoutWrapper\n",
    "\n",
    "from attention import attention as attention\n",
    "from bca_ import *\n",
    "from ordloss import *\n",
    "from utils import *\n",
    "from datautilsbca import *\n",
    "\n",
    "\n",
    "from numpy import array\n",
    "from numpy import argmax\n",
    "from sklearn.preprocessing import LabelEncoder\n",
    "from sklearn.preprocessing import OneHotEncoder\n",
    "from scipy import stats\n",
    "from sklearn.metrics import accuracy_score"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#read data; SEQUENCE_LENGTH is maximum length of sentence in words, SEQUENCE_LENGTH_D is maximum length of document in sentences. \n",
    "SEQUENCE_LENGTH = 40\n",
    "SEQUENCE_LENGTH_D = 25\n",
    "max_vocab = 75000\n",
    "train_split = 0.95\n",
    "BATCH_SIZE = 20\n",
    "\n",
    "# system parameters\n",
    "HIDDEN_SIZE = 150\n",
    "HIDDEN_SIZE_D = 150\n",
    "ATTENTION_SIZE = 75\n",
    "ATTENTION_SIZE_D = 50\n",
    "LAYER_1 = 500\n",
    "LAYER_2 = 250\n",
    "LAYER_3 = 100\n",
    "KEEP_PROB = 0.7\n",
    "#NUM_EPOCHS = 1  # max val_acc at __\n",
    "DELTA = 0.75"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fpath = 'data/TOEFL'\n",
    "\n",
    "#add dict name\n",
    "dict_name = 'bca_dm_toefl_model/dict.csv'\n",
    "# load the dictionary from the pre-trained model folder\n",
    "import csv \n",
    "dictionary = {}\n",
    "for key,val in csv.reader(open(dict_name)):\n",
    "    dictionary[key] = val"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# the test data set; the fformat is csv, with the text column labelled 'text'\n",
    "df_test = pd.read_csv(os.path.join(fpath,'test.csv'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def read_data(raw_text):\n",
    "    content = raw_text\n",
    "    #print(content)\n",
    "    content = content.split() #splits the text by spaces (default split character)\n",
    "    content = np.array(content)\n",
    "    content = np.reshape(content, [-1, ])\n",
    "    return content"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def read_test_set(df_val, dictionary, SEQUENCE_LEN_D = 40, SEQUENCE_LEN = 65, BATCH_SIZE = 10):\n",
    "    \n",
    "    X_val = []\n",
    "    \n",
    "    for i in df_val['text1']:\n",
    "        i = sent_tokenize(i)\n",
    "        #print(i)\n",
    "        X_val.append([dictionary['START_SENT']])        \n",
    "        for j in i[:SEQUENCE_LEN_D-2]:\n",
    "            #print(j)\n",
    "            #print(str(j).lower())\n",
    "            x = read_data(str(j).lower())\n",
    "            #print(x)\n",
    "            data = []\n",
    "            data.append(dictionary['START'])\n",
    "            for word in x:\n",
    "                if word in dictionary:\n",
    "                    index = dictionary[word]\n",
    "                    #count_iv_test += 1\n",
    "\n",
    "                else:\n",
    "                    index = dictionary['UNK']\n",
    "                    #count_oov_test += 1\n",
    "\n",
    "                data.append(index)\n",
    "            data.append(dictionary['END'])\n",
    "            X_val.append(data)\n",
    "        X_val.append([dictionary['END_SENT']])\n",
    "        for k in range(max(SEQUENCE_LEN_D - (len(i)+2), 0)):\n",
    "            X_val.append([0])\n",
    "\n",
    "    print('len of test set: ', len(X_val)//BATCH_SIZE)\n",
    "\n",
    "    rank_val = df_val['label']\n",
    "    target_val = np.array(rank_val)\n",
    "    onehot_encoder = OneHotEncoder(sparse=False)\n",
    "    \n",
    "    integer_encoded = target_val.reshape(len(target_val), 1)\n",
    "    y_test = onehot_encoder.fit_transform(integer_encoded)\n",
    "\n",
    "    return X_val, y_test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X_test, y_test = read_test_set(df_test, dictionary, SEQUENCE_LEN_D = SEQUENCE_LENGTH_D, SEQUENCE_LEN = SEQUENCE_LENGTH)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "doc_vocab_size = len(dictionary)\n",
    "NUM_WORDS = doc_vocab_size\n",
    "EMBEDDING_DIM = 300"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "y_test_len = len(y_test)\n",
    "\n",
    "#use ordinal regression; logistic regression if False\n",
    "ordinal = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def zero_pad_(X, seq_len):\n",
    "    return np.array([x[:seq_len - 1] + [0] * max(seq_len - len(x), 1) for x in X])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Sequences preprocessing\n",
    "vocabulary_size = doc_vocab_size \n",
    "X_test = zero_pad_(X_test, SEQUENCE_LENGTH)\n",
    "\n",
    "#batch size padding \n",
    "X_test = zero_pad_test(X_test, BATCH_SIZE*SEQUENCE_LENGTH_D)\n",
    "y_test = zero_pad_test(y_test, BATCH_SIZE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tf.reset_default_graph()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Different placeholders\n",
    "num_classes = y_test.shape[1]\n",
    "num_classes_s = 8\n",
    "num_classes_s1 = 4\n",
    "batch_ph = tf.placeholder(tf.int32, [None, SEQUENCE_LENGTH])\n",
    "ind_list_ph = tf.placeholder(tf.int32, [None])\n",
    "target_ph = tf.placeholder(tf.float32, [None,num_classes])\n",
    "target_ph_s = tf.placeholder(tf.float32, [None,num_classes_s])\n",
    "target_ph_s1 = tf.placeholder(tf.float32, [None,num_classes_s1])\n",
    "\n",
    "seq_len_ph = tf.placeholder(tf.int32, [None])\n",
    "seq_len_ph_d = tf.placeholder(tf.int32, [None])\n",
    "keep_prob_ph = tf.placeholder(tf.float32)\n",
    "doc_size_ph = tf.placeholder(tf.int32,[None])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Embedding layer\n",
    "embeddings_var = tf.Variable(tf.random_uniform([vocabulary_size, EMBEDDING_DIM], -1.0, 1.0), trainable=True)\n",
    "batch_embedded = tf.nn.embedding_lookup(embeddings_var, batch_ph)\n",
    "batch_embedded = tf.nn.dropout(batch_embedded, keep_prob_ph)\n",
    "\n",
    "W_omega = tf.Variable(tf.random_uniform([HIDDEN_SIZE*2, HIDDEN_SIZE*2], -1.0, 1.0))\n",
    "# (Bi-)RNN layer(-s)\n",
    "with tf.variable_scope('sentence'):\n",
    "    fw_cell = Cell(HIDDEN_SIZE)\n",
    "    bw_cell = Cell(HIDDEN_SIZE)\n",
    "    \n",
    "    fw_cell = DropoutWrapper(fw_cell, input_keep_prob=keep_prob_ph, \n",
    "                             output_keep_prob=keep_prob_ph,state_keep_prob=keep_prob_ph,\n",
    "                             variational_recurrent=True, input_size=batch_embedded.get_shape()[-1], \n",
    "                             dtype = tf.float32)\n",
    "    bw_cell = DropoutWrapper(bw_cell, input_keep_prob=keep_prob_ph, \n",
    "                             output_keep_prob=keep_prob_ph,state_keep_prob= keep_prob_ph,\n",
    "                             variational_recurrent=True, input_size=batch_embedded.get_shape()[-1], \n",
    "                             dtype = tf.float32)\n",
    "    rnn_output, _ = bi_rnn(fw_cell, bw_cell, inputs=batch_embedded, sequence_length=seq_len_ph, dtype=tf.float32)\n",
    "\n",
    "    rnn_outputs = cross_attention(rnn_output, 2 , seq_len_ph, BATCH_SIZE, W_omega, time_major=False, return_alphas=False)\n",
    "    attention_output, alphas = attention(rnn_outputs, ATTENTION_SIZE, seq_len_ph, return_alphas=True)\n",
    "    rnn_outputs_ = cross_attention(rnn_output, SEQUENCE_LENGTH_D, seq_len_ph, BATCH_SIZE, W_omega)\n",
    "    attention_output_, alphas_ = attention(rnn_outputs_ , ATTENTION_SIZE, seq_len_ph, return_alphas = True)\n",
    "    attention_output_ = tf.reshape(attention_output_,[BATCH_SIZE, -1, HIDDEN_SIZE*2*3])\n",
    "    \n",
    "with tf.variable_scope('document'):\n",
    "    fw_cell_d = Cell(HIDDEN_SIZE_D)\n",
    "    bw_cell_d = Cell(HIDDEN_SIZE_D)\n",
    "    \n",
    "    fw_cell_d = DropoutWrapper(fw_cell_d, input_keep_prob=keep_prob_ph, \n",
    "                             output_keep_prob=keep_prob_ph,state_keep_prob=keep_prob_ph,\n",
    "                             variational_recurrent=True, input_size=attention_output_.get_shape()[-1], \n",
    "                             dtype = tf.float32)\n",
    "    bw_cell_d = DropoutWrapper(bw_cell_d, input_keep_prob=keep_prob_ph, \n",
    "                             output_keep_prob=keep_prob_ph,state_keep_prob= keep_prob_ph,\n",
    "                             variational_recurrent=True, input_size=attention_output_.get_shape()[-1], \n",
    "                             dtype = tf.float32)\n",
    "    rnn_outputs_d, _ = bi_rnn(fw_cell_d, bw_cell_d, inputs=attention_output_, \n",
    "                              sequence_length=seq_len_ph_d, dtype=tf.float32)\n",
    "    \n",
    "    #rnn_outputs_d, _ = bi_rnn(Cell(HIDDEN_SIZE_D), Cell(HIDDEN_SIZE_D), inputs=attention_output, sequence_length=seq_len_ph_d, dtype=tf.float32)\n",
    "    attention_output_d, alphas_d = attention(rnn_outputs_d, ATTENTION_SIZE_D, seq_len_ph_d, return_alphas=True)\n",
    "\n",
    "# Dropout\n",
    "drop = tf.nn.dropout(attention_output_d, keep_prob_ph)\n",
    "\n",
    "\n",
    "\n",
    "#first classifier for first task using the representation from attention_outputs\n",
    "#adding more layers... \n",
    "attention_output_sentorder = tf.reshape(attention_output, [BATCH_SIZE, -1])\n",
    "W_s1_ = tf.Variable(tf.truncated_normal([HIDDEN_SIZE*2*2*3, LAYER_1], stddev=0.1))  \n",
    "b_s1_ = tf.Variable(tf.truncated_normal([LAYER_1]))\n",
    "y_hat_s1_ = tf.nn.xw_plus_b(attention_output_sentorder, W_s1_, b_s1_)\n",
    "W_s2 = tf.Variable(tf.truncated_normal([LAYER_1, LAYER_2], stddev=0.1))  \n",
    "b_s2 = tf.Variable(tf.truncated_normal([LAYER_2]))\n",
    "y_hat_s2 = tf.nn.xw_plus_b(y_hat_s1_, W_s2, b_s2)\n",
    "\n",
    "W_s = tf.Variable(tf.truncated_normal([LAYER_2, num_classes_s], stddev=0.1))  \n",
    "b_s = tf.Variable(tf.truncated_normal([num_classes_s]))\n",
    "y_hat_s = tf.nn.xw_plus_b(y_hat_s2, W_s, b_s)\n",
    "y_preds_s = tf.argmax(y_hat_s, axis = 1)\n",
    "loss_s = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(logits=y_hat_s, labels=target_ph_s))\n",
    "\n",
    "#second classifier for second task using the representation from attention_outputs\n",
    "W_s1__ = tf.Variable(tf.truncated_normal([HIDDEN_SIZE*2*2*3, LAYER_1], stddev=0.1))  \n",
    "b_s1__ = tf.Variable(tf.truncated_normal([LAYER_1]))\n",
    "y_hat_s1__ = tf.nn.xw_plus_b(attention_output_sentorder, W_s1__, b_s1__)\n",
    "W_s2_ = tf.Variable(tf.truncated_normal([LAYER_1, LAYER_2], stddev=0.1))  \n",
    "b_s2_ = tf.Variable(tf.truncated_normal([LAYER_2]))\n",
    "y_hat_s2_ = tf.nn.xw_plus_b(y_hat_s1__, W_s2_, b_s2_)\n",
    "\n",
    "W_s1 = tf.Variable(tf.truncated_normal([LAYER_2, num_classes_s1], stddev=0.1))  \n",
    "b_s1 = tf.Variable(tf.truncated_normal([num_classes_s1]))\n",
    "y_hat_s1 = tf.nn.xw_plus_b(y_hat_s2_, W_s1, b_s1)\n",
    "y_preds_s1 = tf.argmax(y_hat_s1, axis = 1)\n",
    "loss_s1 = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(logits=y_hat_s1, labels=target_ph_s1))\n",
    "\n",
    "\n",
    "if ordinal:\n",
    "    # For ordinal regression, same weights for each class\n",
    "    W = tf.Variable(tf.truncated_normal([drop.get_shape()[1].value], stddev=0.1))\n",
    "    W_ = tf.transpose(tf.reshape(tf.tile(W,[num_classes - 1]),[num_classes - 1, drop.get_shape()[1].value]))\n",
    "    b = tf.Variable(tf.cast(tf.range(num_classes - 1), dtype = tf.float32))\n",
    "    y_hat_ = tf.nn.xw_plus_b(drop, tf.negative(W_), b)\n",
    "\n",
    "    # Predicted labels and logits\n",
    "    y_preds, logits = preds(y_hat_,BATCH_SIZE)\n",
    "    y_true = tf.argmax(target_ph, axis = 1)\n",
    "\n",
    "    # Ordinal loss\n",
    "    loss = ordloss_m(y_hat_, target_ph, BATCH_SIZE)\n",
    "    c = stats.spearmanr\n",
    "    str_score = \"Spearman rank:\"\n",
    "\n",
    "else:\n",
    "    W = tf.Variable(tf.truncated_normal([drop.get_shape()[1].value, num_classes], stddev=0.1))  \n",
    "    b = tf.Variable(tf.truncated_normal([num_classes]))\n",
    "    y_hat_ = tf.nn.xw_plus_b(drop, W, b)\n",
    "    # Cross-entropy loss and optimizer initialization\n",
    "    y_preds = tf.argmax(y_hat_, axis = 1)\n",
    "    loss = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(logits=y_hat_, labels=target_ph))\n",
    "    c = accuracy_score\n",
    "    str_score = \"accucary:\"\n",
    "    \n",
    "# Calculate and clip gradients\n",
    "max_gradient_norm = 5\n",
    "lr = 1e-4\n",
    "params = tf.trainable_variables()\n",
    "gradients = tf.gradients(loss, params)\n",
    "clipped_gradients, _ = tf.clip_by_global_norm(gradients, max_gradient_norm)\n",
    "optimizer_ = tf.train.AdamOptimizer(learning_rate=lr)\n",
    "optimizer = optimizer_.apply_gradients(\n",
    "    zip(clipped_gradients, params))\n",
    "\n",
    "#second optimizer for sentence order\n",
    "gradients_s = tf.gradients(loss_s, params)\n",
    "clipped_gradients_s, _ = tf.clip_by_global_norm(gradients_s, max_gradient_norm)\n",
    "optimizer_s = optimizer_.apply_gradients(\n",
    "    zip(clipped_gradients_s, params))\n",
    "\n",
    "#third optimizer for sentence order\n",
    "gradients_s1 = tf.gradients(loss_s1, params)\n",
    "clipped_gradients_s1, _ = tf.clip_by_global_norm(gradients_s1, max_gradient_norm)\n",
    "optimizer_s1 = optimizer_.apply_gradients(\n",
    "    zip(clipped_gradients_s1, params))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "saver = tf.train.Saver()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#testing\n",
    "test_batch_generator = batch_generator(X_test, y_test, BATCH_SIZE, seq_len = SEQUENCE_LENGTH_D, shuffle = False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "MODEL_PATH = \"bca_dm_toefl_model/model300-22275\" "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sess = tf.Session()\n",
    "saver.restore(sess, MODEL_PATH)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#testing on the test set\n",
    "num_batches = X_test.shape[0] // (BATCH_SIZE*SEQUENCE_LENGTH_D)\n",
    "true = []\n",
    "ypreds = []\n",
    "a = []\n",
    "a_d = []\n",
    "true = []\n",
    "preds_ = []\n",
    "log_ = []\n",
    "doc_size_np = np.array([0]*SEQUENCE_LENGTH_D)\n",
    "\n",
    "for bx in range(num_batches):\n",
    "    x_batch, y_batch = next(test_batch_generator)\n",
    "    seq_len = np.array([list(x).index(0) + 1 for x in x_batch])  # actual lengths of sequences\n",
    "    seq_len_d = []               \n",
    "    l = SEQUENCE_LENGTH_D\n",
    "    for i in range(0,len(x_batch),l):\n",
    "        for j in range(i,i+l):\n",
    "            if list(x_batch[j]).index(0) == 0:\n",
    "                seq_len_d.append(j%l)\n",
    "                break\n",
    "            elif j == i+l-1:\n",
    "                seq_len_d.append(l)\n",
    "\n",
    "    seq_len_d = np.array(seq_len_d)\n",
    "\n",
    "    y_preds_, loss_t, alph, alph_d, log = sess.run([y_preds,loss,alphas_,alphas_d, logits],\n",
    "                  feed_dict={batch_ph: x_batch,\n",
    "                        target_ph: y_batch,\n",
    "                        seq_len_ph: seq_len,\n",
    "                        seq_len_ph_d: seq_len_d,\n",
    "                        doc_size_ph: doc_size_np,\n",
    "                        keep_prob_ph: 1.0})\n",
    "    ypreds.extend(y_preds_)\n",
    "    t = np.argmax(y_batch, axis = 1)\n",
    "    true.extend(t)\n",
    "    a.append(alph)\n",
    "    a_d.append(alph_d)\n",
    "    log_.append(log)\n",
    "\n",
    "true = true[:y_test_len]\n",
    "ypreds = ypreds[:y_test_len]\n",
    "\n",
    "spr = c(true, ypreds)\n",
    "\n",
    "if ordinal:\n",
    "    spr = spr[0]\n",
    "print('Test set '+ str_score + str(spr))\n",
    "\n",
    "rank = stats.spearmanr\n",
    "print('sp rho')\n",
    "print(rank(true, ypreds))\n",
    "\n",
    "from sklearn.metrics import cohen_kappa_score as kappa\n",
    "print('qwk')\n",
    "print(kappa(true, ypreds, weights=\"quadratic\"))\n",
    "\n",
    "from scipy.stats import pearsonr\n",
    "print('pearson')\n",
    "print(pearsonr(true,ypreds))\n",
    "\n",
    "print('kappa')\n",
    "print(kappa(true, ypreds, weights=None))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
